package handler

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"
	"myai/core"
	"myai/core/types"
	"myai/service"
	"myai/store/model"
	"myai/utils"
	"net/http"
	"strings"
)

// WebSocket 连接处理handler

type WebSocketHandler struct {
	BaseHandler
	wsService         *service.WebSocketService
	chatOpenAIHandler *ChatOpenAiHandler
	chatHttpHandler   *ChatHttpHandler
}

func NewWebSocketHandler(app *core.AppServer, s *service.WebSocketService,
	db *gorm.DB, chatOpenAIHandler *ChatOpenAiHandler, chatHttpHandler *ChatHttpHandler) *WebSocketHandler {
	return &WebSocketHandler{
		BaseHandler:       BaseHandler{App: app, DB: db},
		wsService:         s,
		chatOpenAIHandler: chatOpenAIHandler,
		chatHttpHandler:   chatHttpHandler,
	}
}

// Client websocket 客户端连接处理
func (h *WebSocketHandler) Client(c *gin.Context) {
	// 获取连接协议
	clientProtocols := c.GetHeader("Sec-WebSocket-Protocol")
	// 协议升级
	ws, err := (&websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		Subprotocols: strings.Split(clientProtocols, ","),
	}).Upgrade(c.Writer, c.Request, nil)

	if err != nil {
		log.Errorf("websocket协议升级失败：%v", err)
		c.Abort()
		return
	}
	// 获取前端传递的客户端Id
	clientId := c.Query("client_id")
	// 把ws连接封装
	client := types.NewWsClient(ws, clientId)
	// 获取用户Id
	userId := h.GetLoginUserId(c)
	if userId == 0 {
		_ = client.Send([]byte("用户Id无效，请重新登录！"))
		c.Abort()
		return
	}
	// 查询用户信息
	var user model.User
	if err := h.DB.Where("id", userId).First(&user).Error; err != nil {
		_ = client.Send([]byte(fmt.Sprintf("用户信息查询失败：%v", err)))
		c.Abort()
		return
	}
	// 把连接保存
	h.wsService.Clients.Put(clientId, client)
	log.Infof("新 websocket 连接成功，IP：%s", c.RemoteIP())
	// 每个ws连接都使用一个协程进行处理
	go func() {
		for {
			// 读取消息
			_, msg, err := client.Receive()
			if err != nil {
				log.Debugf("读取消息出错：%v，连接已经关闭：%s", err, client.Conn.RemoteAddr())
				client.Close()
				// 删除连接
				h.wsService.Clients.Delete(clientId)
				break
			}
			// 定义接收消息对象
			var message types.InputMessage
			// 把接收的消息转换成对象
			err = utils.JsonDecode(string(msg), &message)
			if err != nil {
				log.Errorf("对象转换出错：%v", err)
				continue
			}

			log.Debugf("接收前端传递的消息：%+v", message)
			// 判断消息类型是不是心跳
			if message.Type == types.MsgTypePing {
				// 响应心跳消息
				utils.SendChannelMsg(client, types.ChPing, "pong")
				continue
			}

			// 定义接受聊天消息对象
			var chatMessage types.ChatMessage
			// 从 message.Body 中转换聊天消息对象到 ChatMessage中
			err = utils.JsonDecode(utils.JsonEncode(message.Body), &chatMessage)
			// 判断是不是对话消息，只处理对话消息，其它类型消息全部丢弃
			if err != nil || message.Channel != types.ChChat {
				log.Warnf("【%s】消息类型不进行处理，已经丢弃：%+v", message.Channel, message.Body)
				continue
			}
			// 查询应用角色
			var chatRole model.ChatRole
			err = h.DB.First(&chatRole, chatMessage.RoleId).Error
			// 判断应用角色是否启用
			if err != nil || !chatRole.Enable {
				utils.SendAndFlush(client, "当前聊天角色不存在或者未启用，请更换角色之后再发起对话！！！")
				continue
			}

			// 如果应用角色有模型Id，则使用应用角色的模型Id
			if chatRole.ModelId > 0 {
				chatMessage.ModelId = chatRole.ModelId
			}

			// 查询模型信息
			var chatModel model.ChatModel
			err = h.DB.Where("id", chatMessage.ModelId).First(&chatModel).Error
			if err != nil || chatModel.Enabled == false {
				utils.SendAndFlush(client, "当前AI模型暂未启用，请更换模型后再发起对话！！！")
				continue
			}

			// 定义会话session
			session := &types.ChatSession{
				ClientIP: c.ClientIP(),
				UserId:   userId,
			}

			// 查询会话
			var chat model.ChatItem
			h.DB.Where("chat_id", chatMessage.ChatId).First(&chat)
			if chat.Id > 0 {
				chatModel.Id = chat.ModelId
				chatMessage.RoleId = int(chat.RoleId)
			}

			session.ChatId = chatMessage.ChatId
			session.Tools = chatMessage.Tools
			session.Stream = chatMessage.Stream
			// 复制模型数据
			err = utils.CopyObject(chatModel, &session.Model)
			if err != nil {
				log.Errorf("复制模型数据失败：%v，%v", err, chatModel)
			}

			// 创建可取消的上下文
			ctx, cancel := context.WithCancel(context.Background())
			// 把客户端Id和上下文进行关联
			h.chatOpenAIHandler.SetCtx(clientId, cancel)
			// 发送聊天请求
			err = h.chatOpenAIHandler.Chat(ctx, session, chatRole, chatMessage.Content, client)
			if err != nil {
				log.Errorf("对大模型发送请求失败：%v", err)
				utils.SendAndFlush(client, err.Error())
			} else {
				// 发送消息结束标识
				utils.SendMsg(client, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeEnd})
				log.Infof("回答完毕：%v", message.Body)
			}

			/*// 把客户端Id和上下文进行关联
			h.chatHandler.ReqCancelFunc.Put(clientId, cancel)
			// 发送请求
			err = h.chatHandler.sendMessage(ctx, session, chatRole, chatMessage.Content, client)
			if err != nil {
				log.Errorf("对大模型发送请求失败：%v", err)
				utils.SendAndFlush(client, err.Error())
			} else {
				// 发送消息结束标识
				utils.SendMsg(client, types.ReplyMessage{Channel: types.ChChat, Type: types.MsgTypeEnd})
				log.Infof("回答完毕：%v", message.Body)
			}*/
		}
	}()
}
